package com.baizeli;

import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.reflect.Field;
import java.lang.reflect.Proxy;
import java.util.UUID;

public class ModuleAccess extends ClassLoader
{
	/**
	 * This field is a static method in a bridge class.<br/>
	 * The bridge class is generated by static{} statement.<br/>
	 * This method invokes JavaLangAccess::addOpens(Module,String,Module)
	 * to open a package to another module.<br/>
	 * It looks like:
	 * <blockquote><pre>
	 * public static void export(
	 *     Object module,
	 *     Object packageName,
	 *     Object target)
	 * {
	 *     SharedSecrets.getJavaLangAccess().addExports(
	 *         (Module)module,
	 *         (String)packageName,
	 *         (Module)target
	 *     );
	 * }</pre></blockquote>
	 * This used {@link java.lang.reflect.Proxy} to force access
	 * JavaLangAccess and SharedSecrets.<br/>
	 * This field is null in Java 8.
	 */
	public static final MethodHandle MODULE;
	public static final MethodHandle EXPORT;
	public static final MethodHandle OPEN;
	public static final MethodHandle READ;
	public static final MethodHandles.Lookup LOOKUP;

	public ModuleAccess()
	{
		super(ModuleAccess.class.getClassLoader());
	}

	public Class<?> loading(byte[] code)
	{
		Class<?> clazz = this.defineClass(null, code, 0, code.length);
		try
		{
			Class.forName(clazz.getName(), true, this);
		}
		catch (ClassNotFoundException e)
		{
			e.printStackTrace(System.out);
		}
		return clazz;
	}

	public static Object module(Class<?> clazz)
	{
		if (ModuleAccess.MODULE != null)
		{
			try
			{
				return ModuleAccess.MODULE.invoke(clazz);
			}
			catch (Throwable t)
			{
				exception(t);
			}
		}
		return null;
	}

	public static void export(Object module, Object packageName, Object target)
	{
		if (ModuleAccess.EXPORT != null)
		{
			try
			{
				ModuleAccess.EXPORT.invokeExact(module, packageName, target);
			}
			catch (Throwable t)
			{
				exception(t);
			}
		}
	}

	/**
	 * A wrapper method of {@link ModuleAccess#OPEN} without any exception.
	 * @param module The module of the class you want to access.
	 * @param packageName The package of the class you want to access.
	 * @param target The module of your class
	 */
	public static void open(Object module, Object packageName, Object target)
	{
		if (ModuleAccess.OPEN != null)
		{
			try
			{
				ModuleAccess.OPEN.invokeExact(module, packageName, target);
			}
			catch (Throwable t)
			{
				exception(t);
			}
		}
	}

	public static void read(Object module, Object target)
	{
		if (ModuleAccess.READ != null)
		{
			try
			{
				ModuleAccess.READ.invokeExact(module, target);
			}
			catch (Throwable t)
			{
				exception(t);
			}
		}
	}

	@SuppressWarnings("unchecked")
	public static <T extends Throwable> void exception(Throwable t) throws T
	{
		throw (T) t;
	}

	static
	{
		try
		{
			ModuleAccess access = new ModuleAccess();
			/*
			 * JavaLangAccess class and SharedSecrets class.
			 * In Java 9 and 10, they are in jdk.internal.misc package.
			 * And in higher version, they ar in jdk.internal.access package
			 */
			Class<?> JLA;
			Class<?> shared;
			try
			{
				JLA = Class.forName("jdk.internal.access.JavaLangAccess");
				shared = Class.forName("jdk.internal.access.SharedSecrets");
			}
			catch(ClassNotFoundException ignored)
			{
				JLA = Class.forName("jdk.internal.misc.JavaLangAccess");
				shared = Class.forName("jdk.internal.misc.SharedSecrets");
			}
			// Use proxy to force access the package of JavaLangAccess in the module of the proxy object.
			Object proxy = Proxy.newProxyInstance(access, new Class[]{JLA}, (_1, _2, _3) -> null);
			String packageName = proxy.getClass().getPackage().getName();
			// The name of the bridge class is a random UUID.
			String uuid = 'Z' + UUID.randomUUID().toString().toUpperCase().replaceAll("-", "");
			// In the same package, this class can also access JLA and SS
			String className = packageName.replace('.', '/') + "/" + uuid;
			/*
			 * This bridge class looks like:
			 *
			 * package com.sun.proxy.jdk.proxy1;
			 *
			 * import jdk.internal.access.JavaLangAccess;
			 * import jdk.internal.access.SharedSecrets;
			 *
			 * public class Bridge
			 * {
			 *     public static void export(Object module, Object packageName, Object open)
			 *     {
			 *         SharedSecrets.getJavaLangAccess().addExports((Module)module, (String)packageName, (Module)open);
			 *     }
			 *
			 *     public static void open(Object module, Object packageName, Object open)
			 *     {
			 *         SharedSecrets.getJavaLangAccess().addOpens((Module)module, (String)packageName, (Module)open);
			 *     }
			 *
			 *     public static void read(Object module, Object packageName, Object open)
			 *     {
			 *         SharedSecrets.getJavaLangAccess().addReads((Module)module, (String)packageName, (Module)open);
			 *     }
			 *
			 *     static
			 *     {
			 *         Bridge.class.getModule().addExports(
			 *             "com.sun.proxy.jdk.proxy1",
			 *             Class.forName("org.mve.invoke.ModuleAccess").getModule()
			 *         );
			 *     }
			 * }
			 */
			ClassWriter visitor = new ClassWriter(0);
			visitor.visit(52, Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, className, null, "java/lang/Object", null);
			MethodVisitor mv;

			mv = visitor.visitMethod(
				Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
				"open",
				"(Ljava/lang/Object;Ljava/lang/Object;Ljava/lang/Object;)V",
				null, null
			);
			mv.visitCode();
			mv.visitMethodInsn(
				Opcodes.INVOKESTATIC,
				shared.getTypeName().replace('.', '/'),
				"getJavaLangAccess",
				MethodType.methodType(JLA).toMethodDescriptorString(),
				false
			);
			mv.visitVarInsn(Opcodes.ALOAD, 0);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitVarInsn(Opcodes.ALOAD, 1);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/String");
			mv.visitVarInsn(Opcodes.ALOAD, 2);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitMethodInsn(
				Opcodes.INVOKEINTERFACE,
				JLA.getTypeName().replace('.', '/'),
				"addOpens",
				"(Ljava/lang/Module;Ljava/lang/String;Ljava/lang/Module;)V",
				true
			);
			mv.visitInsn(Opcodes.RETURN);
			mv.visitMaxs(4, 3);
			mv.visitEnd();

			mv = visitor.visitMethod(
				Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
				"export",
				"(Ljava/lang/Object;Ljava/lang/Object;Ljava/lang/Object;)V",
				null, null
			);
			mv.visitCode();
			mv.visitMethodInsn(
				Opcodes.INVOKESTATIC,
				shared.getTypeName().replace('.', '/'),
				"getJavaLangAccess",
				MethodType.methodType(JLA).toMethodDescriptorString(),
				false
			);
			mv.visitVarInsn(Opcodes.ALOAD, 0);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitVarInsn(Opcodes.ALOAD, 1);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/String");
			mv.visitVarInsn(Opcodes.ALOAD, 2);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitMethodInsn(
				Opcodes.INVOKEINTERFACE,
				JLA.getTypeName().replace('.', '/'),
				"addExports",
				"(Ljava/lang/Module;Ljava/lang/String;Ljava/lang/Module;)V",
				true
			);
			mv.visitInsn(Opcodes.RETURN);
			mv.visitMaxs(4, 3);
			mv.visitEnd();

			mv = visitor.visitMethod(
				Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
				"read",
				"(Ljava/lang/Object;Ljava/lang/Object;)V",
				null, null
			);
			mv.visitCode();
			mv.visitMethodInsn(
				Opcodes.INVOKESTATIC,
				shared.getTypeName().replace('.', '/'),
				"getJavaLangAccess",
				MethodType.methodType(JLA).toMethodDescriptorString(),
				false
			);
			mv.visitVarInsn(Opcodes.ALOAD, 0);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitVarInsn(Opcodes.ALOAD, 1);
			mv.visitTypeInsn(Opcodes.CHECKCAST, "java/lang/Module");
			mv.visitMethodInsn(
				Opcodes.INVOKEINTERFACE,
				JLA.getTypeName().replace('.', '/'),
				"addReads",
				"(Ljava/lang/Module;Ljava/lang/Module;)V",
				true
			);
			mv.visitInsn(Opcodes.RETURN);
			mv.visitMaxs(3, 2);
			mv.visitEnd();

			mv = visitor.visitMethod(Opcodes.ACC_STATIC, "<clinit>", "()V", null, null);
			mv.visitCode();
			mv.visitLdcInsn(Type.getType("L" + className + ";"));
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Class", "getModule", "()Ljava/lang/Module;", false);
			mv.visitLdcInsn(packageName);
			mv.visitLdcInsn(ModuleAccess.class.getTypeName());
			mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;", false);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Class", "getModule", "()Ljava/lang/Module;", false);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Module", "addExports", "(Ljava/lang/String;Ljava/lang/Module;)Ljava/lang/Module;", false);
			mv.visitInsn(Opcodes.POP);
			mv.visitLdcInsn(Type.getType("L" + className + ";"));
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Class", "getModule", "()Ljava/lang/Module;", false);
			mv.visitLdcInsn(packageName);
			mv.visitLdcInsn(ModuleAccess.class.getTypeName());
			mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/Class", "forName", "(Ljava/lang/String;)Ljava/lang/Class;", false);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Class", "getModule", "()Ljava/lang/Module;", false);
			mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/lang/Module", "addOpens", "(Ljava/lang/String;Ljava/lang/Module;)Ljava/lang/Module;", false);
			mv.visitInsn(Opcodes.POP);
			mv.visitInsn(Opcodes.RETURN);
			mv.visitMaxs(3, 0);
			mv.visitEnd();

			visitor.visitEnd();
			byte[] code = visitor.toByteArray();

			Class<?> clazz = access.loading(code);
			ModuleAccess.class.getModule().addReads(clazz.getModule());
			// Use MethodHandle to invoke methods.
			MethodHandles.Lookup lookup = MethodHandles.lookup();
			Class<?> moduleClass = Class.forName("java.lang.Module");
			MODULE = lookup.findVirtual(Class.class, "getModule", MethodType.methodType(moduleClass));
			EXPORT = lookup.findStatic(
				clazz,
				"export",
				MethodType.methodType(void.class, Object.class, Object.class, Object.class)
			);
			OPEN = lookup.findStatic(
				clazz,
				"open",
				MethodType.methodType(void.class, Object.class, Object.class, Object.class)
			);
			READ = lookup.findStatic(
				clazz,
				"read",
				MethodType.methodType(void.class, Object.class, Object.class)
			);

			ModuleAccess.open(MethodHandles.Lookup.class.getModule(), MethodHandles.Lookup.class.getPackageName(), ModuleAccess.class.getModule());
			Field implLookupField = MethodHandles.Lookup.class.getDeclaredField("IMPL_LOOKUP");
			implLookupField.setAccessible(true);
			LOOKUP = (MethodHandles.Lookup) implLookupField.get(null);
		}
		catch (Throwable t)
		{
			exception(t);
			throw new RuntimeException(t);
		}
	}
}
